/*
 * Alloy Analyzer
 * Copyright (c) 2002 Massachusetts Institute of Technology
 *
 * The Alloy Analyzer is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * The Alloy Analyzer is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with the Alloy Analyzer; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 */

package alloy.transform;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;

import alloy.ast.ASTDepthFirstReplacer;
import alloy.ast.ComprehensionExpr;
import alloy.ast.Decl;
import alloy.ast.Expr;
import alloy.ast.Let;
import alloy.ast.LetDecl;
import alloy.ast.LetFormula;
import alloy.ast.LetImpl;
import alloy.ast.Node;
import alloy.ast.QuantifiedFormula;
import alloy.ast.SigExpr;
import alloy.ast.SumIntExpr;
import alloy.ast.VarCreator;
import alloy.ast.Variable;
import alloy.ast.VariableExpr;
import alloy.util.Dbg;
import alloy.util.Msg;

/**
   Desugars let construct.
*/

public final class DesugarLetVisitor extends ASTDepthFirstReplacer {

     /** mapping from variable names to expressions which should
        replace them */
    private Map _varToExpr = new HashMap();
	/**
	 * Visit a let node.  Collect the names bound by the let,
	 * and replace their appearances in the body of the let.  
	 * Checks for duplicate names declared in one let.
	 */
    public Object visit(LetImpl letImpl_) {
        return _visitLet(letImpl_);
    }

    public Object visit(LetFormula letForm_) {
        return _visitLet(letForm_);
    }
    
    private Object _visitLet(Let let_) {
    	// keep copy of existing mapping
    	Map oldVarToExpr = new HashMap(_varToExpr);
        // set of names declared in this let
        Set letNames = new HashSet();
        // iterate through declarations
        Iterator letDeclIter = let_.getLetDecls().getLetDeclIter();
        while (letDeclIter.hasNext()) {
	    
            LetDecl letDecl = (LetDecl)letDeclIter.next();
	    letDecl.applyReturnVisitor(this);
     
	    String name = letDecl.getVar().nodeString();
            if (letNames.contains(name)) { // duplicate name
                Dbg.user(new Msg("duplicate name declaration in let", let_));
                continue;
            }

            // add binding to mapping, will override any existing mapping for name
            _varToExpr.put(name, letDecl.getExpr());
        }
        Node body = let_.getBody();
        // visit body and save return value (since it may itself be a let)
        body = (Node)body.applyReturnVisitor(this);
        body.annotateTransform("desugaring of let");
    	// restore original mapping
    	_varToExpr = oldVarToExpr;
        // body replaces let	
       return body;		
    }        
	/**
	 * Visit a var. creator.  Ensure that names bound by
	 * the var. creator are not replaced in its body according
	 * to an outer let binding.
	 */
    private VarCreator _handleVarCreator(VarCreator varCreator_) {
        // keep a copy of existing mapping
        Map oldVarToExpr = new HashMap(_varToExpr);
        Iterator declIter = varCreator_.getDecls().getDeclIter();
        while (declIter.hasNext()) {
            Decl decl = (Decl)declIter.next();
            // visit declaration
            decl.applyReturnVisitor(this);
            // remove decl's names from mapping
            _removeDeclNames(decl);
        }
        // visit body
        varCreator_.setBody((Node)varCreator_.getBody().applyReturnVisitor(this));
        // restore original mapping
        _varToExpr = oldVarToExpr;
        return varCreator_;
    }

    public Object visit(QuantifiedFormula quantForm_) {
    	return _handleVarCreator(quantForm_);
    }

    public Object visit(ComprehensionExpr compExpr_) {
    	return _handleVarCreator(compExpr_);
    }

    public Object visit(SumIntExpr sumIntExpr_) {
    	return _handleVarCreator(sumIntExpr_);
    }
    
    /** replace variable with expression it is bound to if necessary */
    public Object visit(VariableExpr varExpr_) {
    	Expr ret = varExpr_;
        String varString = varExpr_.nodeString();
        if (_varToExpr.containsKey(varString)) {
            ret = (Expr)((Expr)_varToExpr.get(varString)).copy();
            ret.annotateTransformReplacing("desugaring of let", varExpr_);
        }
       return ret;
    }
    
    public Object visit(SigExpr sigExpr_) {
    	Expr ret = sigExpr_;
    	String varString = sigExpr_.nodeString();
        if (_varToExpr.containsKey(varString)) {
            ret = (Expr)((Expr)_varToExpr.get(varString)).copy();
            ret.annotateTransformReplacing("desugaring of let", sigExpr_);
        }
       return ret;
    }
    
    private void _removeDeclNames(Decl decl) {
        Iterator varIter = decl.getVariables().getVariableIter();
        while (varIter.hasNext()) {
            Variable var = (Variable)varIter.next();
            String name = var.nodeString();
            if (_varToExpr.containsKey(name)) {
                _varToExpr.remove(name);
            }
        }
    }

}


